import atexit
import logging
from multiprocessing import Process, Manager
from threading import Thread, Event
from websocket import WebSocketApp
from utils.config import ConnectionConfig
from utils.exceptions import ConnectionError
import websocket
import threading
import json
import time
from queue import Queue

logger = logging.getLogger(__name__)

class BaseClient:
    """WebSocket客户端基类"""
    def __init__(self, config: ConnectionConfig, spi=None):
        """
        初始化基础客户端
        
        Args:
            config: 连接配置
            spi: 回调接口实现
        """
        self.config = config
        self.spi = spi
        self.ws = None
        self.connected = False
        self.ws_thread = None
        self.task_queue = Queue()  # 增加消息队列
        self._init_components()
        
    def _init_components(self):
        """初始化核心组件"""
        self._data_queue = Manager().Queue()
        self._connection_event = Event()
        self._init_ws()
        self._register_cleanup()
        if self.spi:
            self.spi.task_queue = self.task_queue  # 注入队列到SPI
            # 启动消息处理线程
            self.process_thread = threading.Thread(
                target=self.spi.start, 
                args=(self.task_queue,),
                daemon=True
            )
            self.process_thread.start()

    def _init_ws(self):
        """初始化WebSocket连接"""
        #websocket.enableTrace(True)  # 启用WebSocket库的详细跟踪
        self.ws = WebSocketApp(
            url=self._get_ws_url(),
            on_open=self._on_open,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close
        )

    def _get_ws_url(self) -> str:
        """获取WebSocket地址（由子类实现）"""
        raise NotImplementedError

    def connect(self):
        """建立连接"""
        if self.connected:
            return
            
        url = self._get_ws_url()
        logger.info(f"正在连接到 {url}")
        
        # 确保回调函数被正确设置
        #websocket.enableTrace(True)  # 启用WebSocket库的详细跟踪
        self.ws = websocket.WebSocketApp(
            url,
            on_open=self._on_open,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close
        )
        
        self.ws_thread = threading.Thread(target=self.ws.run_forever)
        self.ws_thread.daemon = True
        self.ws_thread.start()
        
        # 等待连接建立
        if not self._wait_connection():
            logger.error("连接超时")
            self.close()
            return False
        
        logger.info("连接建立成功,建立回报线程")
        # 启动消息处理线程
        self._start_processing()
        
        return True

    def _wait_connection(self, timeout=15):
        """
        等待连接建立
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            bool: 连接是否成功建立
        """
        start_time = time.time()
        while not self.connected:
            if time.time() - start_time > timeout:
                return False
            time.sleep(0.1)
        
        return True

    def _start_processing(self):
        """启动处理线程"""
        self._process_thread = Thread(target=self._process_messages, daemon=True)
        self._process_thread.start()

    def _process_messages(self):
        """消息处理循环"""
        while True:
            data = self._data_queue.get()
            if data is None:
                break
            try:
                self._handle_message(data)
            except Exception as e:
                logger.error(f"消息处理失败: {str(e)}")

    def _handle_message(self, message):
        """处理接收到的消息"""
        try:
            # 尝试解析JSON消息
            if isinstance(message, str):
                data = json.loads(message)
            else:
                # 如果已经是字典，直接使用
                data = message
            # 调用SPI回调
            if self.spi:
                try:
                    self.spi.on_message(message)
                except Exception as e:
                    logger.error(f"SPI回调失败: {str(e)}", exc_info=True)
        except json.JSONDecodeError:
            logger.warning(f"收到非JSON格式消息: {message[:100]}...")
        except Exception as e:
            logger.error(f"处理消息时发生未知错误: {e}", exc_info=True)

    def _on_open(self, ws):
        """连接建立回调"""
        logger.info("连接已建立")
        self.connected = True
        if self.spi:
            self.spi.on_connected()
        self._connection_event.set()

    def _on_message(self, ws, message):
        """WebSocket消息回调"""
        logger.debug(f"收到原始消息: {message}")
        try:
            if self.task_queue:
                self.task_queue.put(message)  # 将消息放入队列
        except Exception as e:
            logger.error(f"消息入队失败: {str(e)}")

    def _on_error(self, ws, error):
        """错误处理回调"""
        logger.error(f"连接错误: {str(error)}")
        self.connected = False
        if self.spi:
            self.spi.on_error(error)
        self._connection_event.clear()

    def _on_close(self, ws, close_status_code, close_msg):
        """连接关闭回调"""
        logger.info(f"连接关闭: {close_status_code} - {close_msg}")
        self.connected = False
        if self.spi:
            self.spi.on_disconnected()
        self._connection_event.clear()

    def _register_cleanup(self):
        """注册清理函数"""
        atexit.register(self.close)

    def close(self):
        """关闭连接"""
        if self.ws:
            self.ws.close()
        self._data_queue.put(None)
        if self.ws_thread and self.ws_thread.is_alive():
            self.ws_thread.join(timeout=1.0)
        logger.info("连接已安全关闭") 